#!/usr/bin/env python

# This file is part of GMultiserver
# Copyright (C) 2012 Gustaw Lewandowski 
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307 USA

from utils import config, GServer, GClient
from mime import get_header, get_param, get_section, get_bodystructure, get_envelope, convert_to_b64, convert_to_binary, structure, convert_to_log, save_message, message_from_string
from imaplib import CRLF, Int2AP
import sqlite3
from utils import GAnswer
import cPickle as pickle

def text_part(s):
	s = str(s)
	h = s.find(2 * CRLF)
	if h == -1:
		h = s.find('\n\n')
		if h != -1:
			s = s[h + 1:]
	else:
		s = s[h + 2:]
	return s

class Imap4Client(GClient):
	tagnum = 0
	tagpre = None
	def tag_init(self):
		if self.tagpre is None:
			import random, re
			self.tagpre = Int2AP(random.randint(4096, 65535))
			self.tagre = re.compile(r'(?P<tag>' + self.tagpre + r'\d+) (?P<type>[A-Z]+) (?P<data>.*)')

	def tag(self):
		self.tagnum = self.tagnum + 1
		return '%s%s' % (self.tagpre, self.tagnum)

	def send_cmd(self, cmd, data = None, addTag = True, fetch = False, expect_msg = False):
		self.tag_init()	
		tag = self.tag()
		if addTag:
			cmd = '%s %s' % (tag, cmd)

		ans = ''
		status = ''
		load = ''
		msgs = {}
		s = ''

		a = self.send('%s%s' % (cmd, CRLF))
		if not a.ok: return a
		self.log('sent: %s' % cmd)
		while True:
			a = self.get_line()
			if not a.ok: return a
			j = a.str
			if fetch and j.find('{') != -1:
				size = int(j[j.find('{') + 1:j.rfind('}')])
				self.log('reading %d bytes' % size)
				a = self.recv(size)
				if not a.ok: return a
				if expect_msg:
					msgs[j.split()[1]] = a.str
				else:
					s += j + a.str
			else:
				self.log_debug('get: %s' % j[:-2])
				if j[:len(tag)] == tag:
					status = j.split()[1]
					ans = j[j.find(status):-2]
					break
				elif j[0] == '+' and data is not None:
					self.log('sending %i bytes of data' % len(data))
					self.send_chunking(data)
					self.send(CRLF)
				elif j[0] == '+':
					status = 'OK'
					ans = j[:-2]
					break
				elif j != ')\r\n':
					s += j
		return GAnswer(status = status, str = s, load = ans, msgs = msgs)

	def fetch(self, what, save = False, struct = False, expect_msg = False, cmd = 'FETCH'):
		ans = self.send_cmd('%s %s' % (cmd, what), fetch = True, expect_msg = expect_msg)
		for i in ans.msgs:
			try_add_boudary = False
			if what.find('TEXT') != -1:
				try_add_boudary = True
			m = message_from_string(ans.msgs[i], try_add_boudary)
			if save:
				import time
				save_message(m, '%s/%s/%s/%s' % (config.main.work_dir, self.name, str(time.time()), i))
			if struct:
				ans.struct += 'struct:\n' + structure(m)
			convert_to_log(m)
			ans.load += '\n' + m.as_string()
		return ans

	def login(self, login, password):
		return self.send_cmd('LOGIN %s %s' % (login, password))

	def select(self, folder):
		return self.send_cmd('SELECT %s' % (folder))

	def auth_digest_md5(self, login, password):
		self.tag_init()
		ans = self.send_cmd('AUTHENTICATE DIGEST-MD5')
		from digest_md5 import generate_response
		res = generate_response('imap/', login, password, ans.str)
		ans = self.send_cmd(res, addTag = False)
		self.send(CRLF)
		import base64
		ans.str = base64.b64decode(ans.str)
		return ans

	def append(self, msg, folder = 'INBOX', flags = '', add_date = False, binary = False):
		m = config.get_obj(msg)
		if m is None: return GAnswer(False, load = 'no message %s' % msg)
		tild = ''
		if binary: tild = '~'
		payload = m.as_string(binary = binary)
		d = ''
		if add_date and m.get_param('data') != 'NIL': d = m.get_param('data') + ' '
		if flags != '':
			return self.send_cmd('APPEND %s %s(%s) %s{%i}' % (folder, d, flags, tild, len(payload)), data = payload)
		return self.send_cmd('APPEND %s %s%s{%i}' % (folder, d, tild, len(payload)), data = payload)

	def setmetadata(self, path = '', folder = 'inbox', msg = None):
		if msg is not None:
			m = config.get_obj(msg)
			if m is None: return GAnswer(False, load = 'no message %s' % msg)
			payload = m.as_string()
			return self.send_cmd('setmetadata "%s" (%s {%i}' % (folder, path, len(payload)), data = payload)
		return self.send_cmd('setmetadata "%s" (%s)' % (folder, path))

	def getmetadata(self, path = '', folder = 'inbox', expect_msg = False, save = False, struct = False):
		if expect_msg:
			return self.fetch(what = '%s %s' % (folder, path), save = save, struct = struct, cmd = 'getmetadata')
		return self.send_cmd('getmetadata %s %s' % (folder, path))


# Imap4 Server with sqlite3 base
class Imap4Server(GServer):

	def base_init(self):
		self.log('base_init')
		self.db = sqlite3.connect(":memory:")
		self.next_user_id = 0
		c = self.db.cursor()
		c.execute('create table mail (user int, uid int, date text, flags text, msg blob, folder text)')
		c.execute('create table user (user int, nextuid int, login text, pass text)')
		self.db.commit()
		c.close()
		return GAnswer(True, 'DB initiated')

	def populate(self, cfg_name):
		self.log('populate %s' % cfg_name)
		c = config.get(cfg_name)
		m = None
		load = 'populate:\n'
		for e in c:
			if 'seq' in e:
				seq = config.get_obj(e['seq'])
				data_list = []
				if 'data' in e: data_list = e['data']
				for x in seq.gen():
					login = e['login'].replace(seq.var, str(x))
					self.add_mailbox(login, e['password'])
					for d in data_list:
						m = config.get_obj(d['message'])
						if m is None: return GAnswer(False, load = 'no message %s' % msg)
						self.append(login, d['folder'], d['date'], d['flags'], m)
				load += '%s : %s\n' % (repr(seq), data_list)
		return GAnswer(True, load[:-1])

	def append(self, user, folder, date, flags, m, user_id = None, next_uid = None):
		if not self.folder_names_case_sensitive:
			folder = folder.lower()
		c = self.db.cursor()
		if user_id is None or next_uid is None:
			r = c.execute('SELECT * FROM user WHERE login=:u;', {'u': user})
			user_id = 0
			next_uid = 0
			for i in r:
				user_id = int(i[0])
				next_uid = int(i[1])
		# todo not found user
		c.execute('INSERT INTO mail VALUES (?,?,?,?,?,?);', (user_id, next_uid, date, flags, sqlite3.Binary(pickle.dumps(m)), folder))
		next_uid += 1
		c.execute('UPDATE user SET nextuid=:next_uid WHERE user=:user;', {'next_uid': next_uid, 'user': user_id})
		c.close()
		self.db.commit()
		self.log('append message user:%s uid:%d folder:%s' % (user, next_uid, folder))

	def add_mailbox(self, user, passwd):
		self.log('add_mailbox login:%s' % user)
		c = self.db.cursor()
		c.execute('INSERT INTO user VALUES (?,?,?,?);', (self.next_user_id, 1 , user , passwd))
		c.close()
		self.db.commit()
		self.next_user_id += 1
	
	def handler(self, socket, address):
		self.send_line(socket, '* OK [CAPABILITY IMAP4rev1] GMultiServer IMAP4rev1 Ready')
		next_uid = 1
		user_id = -1
		user = ''
		folder = 'INBOX'
		uid_val = '1234'
		while True:
			try:
				r = socket.recv(256)
			except Exception, e:
				self.log_error(e)
				socket.close()
				return None
			if r == '':
				socket.close()
				return None
			self.log('get: ' +r[:-1])
			a = r[:-2].split(' ')
			if a[1].upper() == 'LOGIN':
				a[2] = a[2].replace('"', '')
				a[3] = a[3].replace('"', '')
				c = self.db.cursor()
				r = c.execute('SELECT * FROM user WHERE login=:u;', {'u': a[2]})
				auth_ok = False
				rs = self.login_resp_fail
				for i in r:
					user_id = int(i[0])
					next_uid = int(i[1])
	
					if i[3] == a[3]:
						auth_ok = True
						rs = self.login_resp_ok
						self.log('password match')
					else:
						self.log('wrong password')
					break

				if user_id != -1:
					self.log('user_id:%d, next_uid:%d' % (user_id, next_uid))

				self.send_line(socket, '%s %s' % (a[0], rs))
				user = a[2]
				c.close()
				if auth_ok:
					break
			elif a[1].upper() == 'CAPABILITY':
				self.send_line(socket, '* CAPABILITY %s' % self.capability_str)
				socket.send(a[0] + ' OK CAPABILITY completed\r\n')
		uids = {}
		while True:
			self.log('waiting for cmd')
			uid_cmd = False
			try:
				r = socket.recv(1024)
			except Exception, e:
				self.log(e)
				break
			if r == '':
				break
			a = r[:-2].split(' ')
			self.log('get: ' +r[:-1])
			command = a[1].upper()
			if command == 'UID': # map UID command to non UID eqv
				uid_cmd = True
				command = a[2].upper()
				a = [a[0]] + a[2:]

			if command == 'NOOP' or command == 'CREATE' or command == 'SUBSCRIBE':
				self.send_line(socket, a[0] + ' OK %s completed' % command)
			elif command == 'NAMESPACE':
				self.send_line(socket, '* NAMESPACE (("" "/")) NIL NIL')
				self.send_line(socket, a[0] + ' OK Success')
			elif command == 'GETMETADATA':
				# only for clint  testing
				self.send_line(socket, '* METADATA "INBOX" (/private/vendor/value "test")')
				self.send_line(socket, a[0] + ' OK GETMETADATA completed')
			elif command == 'LIST' or command == 'LSUB':
				#folder = a[3][1:-1]
				folder = 'INBOX'
				if a[2] == '""' and a[3] == '""':
					self.send_line(socket, '* %s (\\Noselect) "/" "/"' % command)
				else:
					self.send_line(socket, '* %s (\HasNoChildren) "/" "INBOX"' % command)
				self.send_line(socket, a[0] + ' OK %s completed' % command)
			elif command == 'LOGOUT':
				self.send_line(socket, '* BYE Nice talking to you')
				self.send_line(socket, a[0] + ' OK LOGOUT successful')
				socket.close()
				break
			elif command == 'SEARCH':
				ml = self.get_message_list(user_id, '1:*', folder)
				if a[2] == 'UNSEEN':
					s = ''
					for i in ml:
						if i[3].upper().find('\\SEEN') == -1:
							s += str(i[1]) + ' '
					s = s[:-1]
				if a[2] == 'ALL':
					s = ''
					for i in ml:
						s += str(i[1]) + ' '
					s = s[:-1]
				self.send_line(socket, '* SEARCH ' + s)
				self.send_line(socket, a[0] + ' OK SEARCH completed')
			elif command == 'STATUS':
				slist = set(r[r.find('(')+1:r.find(')')].split(' '))
				ml = self.get_message_list(user_id, '1:*', folder)
				s = ' ('
				v = ''
				for i in slist:
					if i.upper() == 'UIDNEXT':
						v = str(next_uid)
					elif i.upper() == 'UIDVALIDITY':
						v = uid_val
					elif i == 'MESSAGES':
						v = str(len(ml))
					elif i == 'UNSEEN':
						cnt = 0
						for j in ml:
							if j[3].upper().find('\\SEEN') == -1:
								cnt += 1
						v = str(cnt)
					s += i + ' ' + v + ' '
				s = s[:-1] + ')'
				self.send_line(socket, '* STATUS ' + a[2] + s)
				self.send_line(socket, a[0] + ' OK STATUS completed')
			elif command == 'CAPABILITY':
				self.send_line(socket, '* CAPABILITY %s' % self.capability_str)
				socket.send(a[0] + ' OK CAPABILITY completed\r\n')
			elif command == 'EXPUNGE':
				# delete msgs
				self.send_line(socket, a[0] + ' OK Completed')
				ml = self.get_message_list(user_id, '1:*', folder)
				if not self.read_only:
					c = self.db.cursor()
					for i in ml:
						if i[3].upper().find('\\DELETED') != -1:
							self.send_line(socket,'* %d EXPUNGE' % i[1])
							c.execute('DELETE FROM mail WHERE user=:user AND uid=:uid;', {'uid': i[1], 'user': user_id})
					c.close()
					self.db.commit()
			elif command == 'STORE':
				ml = self.get_message_list(user_id, a[2], folder, uids, uid_cmd)
				new_flags = set(r[r.find('(')+1:r.find(')')].split(' ')) # remove braces, split and conv to set
				c = self.db.cursor()
				for i in ml:
					flags = set(i[3].split())
					if a[3][0] == '+':
						flags = flags.union(new_flags)
					else:
						flags = flags.difference(new_flags)
					# convert set to string avoid '\\\\'
					s = ''
					for j in flags:
						if j == '':
							continue
						s += chr(92) + j.replace(chr(92),'') + ' '
					s = s[:-1]
					if not self.read_only:
						c.execute('UPDATE mail SET flags=:flags WHERE uid=:uid AND user=:user;', {'uid': i[1], 'user': user_id, 'flags': s})
					
					if a[3][-6:].upper() != 'SILENT':
						self.send_line(socket,'* %d FETCH (FLAGS (%s))' % (i[1], s))
				c.close()
				self.db.commit()
				self.send_line(socket, a[0] + ' OK STORE completed')
			elif command == 'SELECT':
				folder = a[2]
				uids = {}
				if folder[0] == '"':
					folder = folder[1:-1]
				if not self.folder_names_case_sensitive:
					folder = folder.lower()
				ml = self.get_message_list(user_id, '1:*', folder)
				seen_cnt = 0
				j = 1
				uid_map = {}
				for i in ml:
					uids[i[1]] = j
					uid_map[j] = i[1]
					j += 1
					if i[3].upper().find('\\SEEN') != -1:
						seen_cnt += 1
				self.send_line(socket, '* %d EXISTS\r\n* %d RECENT\r\n* FLAGS (\Seen \Unseen \Deleted \Flagged \Answered \Recent)\r\n* OK [UIDVALIDITY 1234]' % (len(ml), len(ml) - seen_cnt))
				self.send_line(socket, a[0] + ' OK [READ-WRITE] SELECT successful')
				self.log('UID mapping: %s' % uid_map)
			elif command == 'APPEND':
				if self.append_enable:
					size = int(a[-1][a[-1].find('~')+2:-1])
					date = r[r.find('"')+1:r.rfind('"')]
					flags = ''
					if r.find('(') != -1:
						flags  = r[r.find('(')+1: r.find(')')]
					self.send_line(socket, '+ Ready for %d octets of data' % size)
					cnt = size / 8192
					r = ''
					r = self.get_chunkinig(size, socket)
					m = message_from_string(r)
					convert_to_binary(m)
					socket.recv(2)
					self.append(user, folder, date, flags, m, user_id, next_uid)
					# we need to update uids map
					ml = self.get_message_list(user_id, '1:*', folder)
					uid_map = {}
					j = 1
					for i in ml:
						uids[i[1]] = j
						uid_map[j] = i[1]
						j += 1
					self.log('UID mapping: %s' % uid_map)
					self.send_line(socket, '* OK APPEND Completed')
					self.send_line(socket, a[0] + ' OK [APPENDUID %d %d] Completed' % (1234, next_uid))
					next_uid += 1
				else:
					self.send_line(socket, '%s %s' % (a[0], self.append_resp_fail))
			elif command == 'FETCH':
				if r.find('"') != -1:
					flist = r[r.find('"')+1:r.rfind('"')].split(' ')
				elif r.find('(') != -1:
					flist = r[r.find('(') + 1:r.rfind(')')].split(' ')
				else:
					flist = a[3:]
				ml = self.get_message_list(user_id, a[2], folder, uids, uid_cmd)
				for i in ml:
					#m = message_from_string(i[4])
					m = pickle.loads(str(i[4]))
					resp = '* %d FETCH (' % i[1]
					for q in flist:
						q = q.upper()
						if q == 'FLAGS':
							resp += 'FLAGS (' + i[3] + ') '
						elif q == 'INTERNALDATE':
							resp += 'INTERNALDATE ' + i[2] + ' '
						elif q == 'RFC822.SIZE':
							resp += 'RFC822.SIZE ' + str(len(i[4])) + ' '
						elif q == 'RFC822.HEADER':
							s = get_header(m)
							resp += 'RFC822.HEADER {%d}%s%s%s ' % (len(s) + 2, CRLF, s, CRLF)
						elif q == 'UID':
							resp += 'UID %d ' % i[1]
						elif q == 'BODYSTRUCTURE':
							resp += 'BODYSTRUCTURE %s ' % get_bodystructure(m)
						elif q == 'ENVELOPE':
							resp += 'ENVELOPE %s ' % get_envelope(m)
						elif q == 'RFC822.PEEK' or q == 'RFC822':
							convert_to_b64(m)
							s = m.as_string()
							self.send_line(socket, '%s%s {%d}' % (resp, cmd, len(s)))
							self.send_chunking(s, socket)		
						elif q[:5] == 'BODY[' or q[:10] == 'BODY.PEEK[' or q[:7] == 'BINARY[' or q[:12] == 'BINARY.PEEK[':
							if q[:6] == 'BINARY':
								cmd = 'BINARY'
							else:
								cmd = 'BODY'
							q = r[r.find('[') + 1:r.rfind(']')]
							if q.find('(') != -1:
								q = q[:q.find('(')] + q[q.find(')') + 1:]
							if q == '' or q == 'TEXT':
								if cmd == 'BINARY':
									s = i[4]
									if q == 'TEXT':
										s = text_part(s)
									self.send_line(socket, '%s%s[] ~{%d}' % (resp, cmd, len(s)))
									self.send_chunking(s, socket)
								else:
									convert_to_b64(m)
									s = m.as_string()
									if q == 'TEXT':
										s = text_part(s)
									self.send_line(socket, '%s%s[%s] {%d}' % (resp, cmd, q, len(s)))
									self.send_chunking(s, socket)
								resp = ''
							else:
								secs = q.split(' ')
								for sec in secs:
									if sec == '':
										continue
									self.log('body section: %s' % sec)
									if sec == 'HEADER':
										s = get_header(m)
										resp += '%s[HEADER] {%d}%s%s%s ' % (cmd, len(s) + 2, CRLF, s, CRLF)
									elif sec == 'HEADER.FIELDS':
										q1 = r.find('HEADER.FIELDS')
										fields = r[r.find('(', q1) + 1:r.find(')', q1)].split(' ')
										found_fields = ''
										s = ''
										for f in fields:
											par = get_param(m, f, False)
											if par != 'NIL':
												found_fields += f + ' '
												s += '%s: %s\r\n' % (f, par)
										found_fields = found_fields[:-1]
										if len(s) > 0:
											resp += '%s[HEADER.FIELDS (%s)] {%d}\r\n%s ' % (cmd, found_fields, len(s), s)
									else:
										# e.g. sectiont 1.1.1
										ms = get_section(sec, m)
										if ms is None:
											self.log_error('wrong section %s' % sec)
											break
										if cmd == 'BODY':
											convert_to_b64(ms)
										s = ms.as_string()
										s = s[s.find('\n\n')+2:]
										size = len(s)
										self.send_line(socket, '%s%s[%s] {%d}' % (resp, cmd, sec, size))
										self.send_chunking(s, socket)
										resp = ''
					self.send_line(socket, resp[:-1] + ')')
				self.send_line(socket, a[0] + ' OK FETCH completed')
			else:
				self.send_line(socket, a[0] + ' BAD Function not implemented')

	def get_message_list(self, user_id, uid, folder = 'INBOX', uids = {}, by_uid = True):
		c = self.db.cursor()
		q = c.execute("SELECT * FROM mail WHERE user=:u AND folder=:f;", {'u': user_id, 'f': folder})
		ml = []
		ulist = []
		u = uid.split(':')
		if len(u) == 1:
			u = uid.split(',')
			if len(u) > 1:
				ulist = u
			else:
				ulist = [uid]

		a = int(u[0])
		b = a
		if len(u) == 2:
			if u[1] == '*':
				b = 1000000
			else:
				b = int(u[1])
		for i in q:
			u = 1
			if by_uid:
				u = i[1]
			else:
				if i[1] in uids:
				    u = uids[i[1]]
				else:
				    continue
			if u in ulist:
				ml.append(i)
			elif u >= a and u <= b:
				ml.append(i)
		c.close()
		return ml

	def send_line(self, socket, line):
		self.log('send:' + line)
		try:
			socket.send(line + CRLF)
		except Exception, e:
			self.log_error(e)
